#!/usr/bin/python 
import sys
import os
import random
import math
import copy
from automaton import *


DOT_BIN = "dot"
NEATO_BIN = "neato"
STATE_PREFIX = "Q"
FAILURE_TRANSITION = "x"
GROUND_STATE_NAME = "Garbage"
RNA_EDGE_SUFFIX = ""
FORBIDDEN_SUFFIX = "_F"
MANDATORY_SUFFIX = "_M"
INITIAL_PREFIX = "e_"

def getSucc(q):
  result = None
  trans = q.getOutTransition()
  for t in trans:
    if (t.getLabel()!=FAILURE_TRANSITION) and not(t.getLabel().endswith(RNA_EDGE_SUFFIX)):
      if result is None:
        result = t.getSucc()
      else:
        return None
  return result

def getSuccessors(q):
  result = []
  trans = q.getOutTransition()
  for t in trans:
    if (t.getLabel()!=FAILURE_TRANSITION) and not(t.getLabel().endswith(RNA_EDGE_SUFFIX)):
      result.append(t.getSucc())
  return result



def mergeLists(res, q1,q2):
  res[q1].extend(res[q2])
  for q in res[q2]:
    res[q] = res[q1]

def getTree(a):
  res = {}
  for q in a.getStates():
    if q.getName() != GROUND_STATE_NAME:
      if q not in res:
        res[q] = [q]
      s = getSucc(q)
      if (s is not None) and (s.getName() != GROUND_STATE_NAME):
        if s not in res:
          res[s] = [s]
        mergeLists(res,q,s)
  finalRes = {}
  for q in a.getStates():
    if (q.getName() != GROUND_STATE_NAME):
      if res[q][0] == q:
        finalRes[q] = res[q]
  g = a.getState(GROUND_STATE_NAME)
  if g is not None:
    l = [g]
    finalRes[g] = l
    res[g] = l
  return (finalRes,res)

def getClusters(a):
  res = {}
  tmp = {}
  for q in a.getStates():
    lbl = getListIn(q)
    if lbl not in tmp:
      tmp[lbl] = []
    tmp[lbl].append(q)
    res[q] = tmp[lbl]
    
  finalRes = {}
  for q in a.getStates():
    if (q.getName() != GROUND_STATE_NAME):
      if res[q][0] == q:
        finalRes[q] = res[q]
  return (finalRes,res)



NORMAL_STYLE = "shape = ellipse, style=solid, color=black"
FORBIDDEN_STYLE = "shape = doublecircle, style=solid, color=black"
MANDATORY_STYLE = "shape = doubleoctagon, style=solid, color=black"
FINAL_STYLE = "shape = doublecircle, style=solid, color=black"


def drawAutomaton(a,outpath = "out.ps"):
  newLabels = {}
  newName = {}
  currLabel = 1
  initState = None
  for q in a.getStates():
    if q.isInitial():
      initState = q
    newName[q.getName()] = STATE_PREFIX+q.getName()
    newLabels[q.getName()] = q.getName()
    offset = newLabels[q.getName()].find("_")
    if (offset!=-1) and not (q.getName().startswith(INITIAL_PREFIX)):
      newLabels[q.getName()] = newLabels[q.getName()][:offset]
    currLabel += 1

  # Extracting clusters
  (tree,hs) = getClusters(a)
  #print tree

  # Helping Gviz overcome its own stupidity by partially precomputing ranks
  initNodes = set()
  ranks = {0:[]}
  for q in a.getStates():
    if q.getName().startswith(INITIAL_PREFIX):
      initNodes.add(q)
      ranks[0].append(q)
  depth = 0
  while depth in ranks:
    for q in ranks[depth]:
      if q.getName() != GROUND_STATE_NAME:
        l = getSuccessors(q)
        for s in l:
          if (s.getName() != GROUND_STATE_NAME):
            if depth+1 not in ranks:
              ranks[depth+1] = []
            ranks[depth+1].append(s)
    depth += 1
  ranktxt = ""
  for depth in ranks:
    tmp = ""
    for q in ranks[depth]:
      tmp += "%s; "%(newName[q.getName()])    
    ranktxt += "\n {     rank=same; %s}" %(tmp)
  #ranktxt = ""
  

  # Edges list
  edgesTxt = ""
  for q in a.getStates():
    for t in q.getOutTransition():
      qp = t.getSucc()
      lbl = t.getLabel()
#      if (hs[qp] != hs[q]) or (lbl==FAILURE_TRANSITION) or (lbl.endswith(RNA_EDGE_SUFFIX)):
      if (1):
        style = ""
        switch = 0
        if ((q.getName()==GROUND_STATE_NAME)or (qp.getName()==GROUND_STATE_NAME)):
          style = ", style = dashed"
          switch = 1
        elif ((lbl==FAILURE_TRANSITION) or (lbl.endswith(RNA_EDGE_SUFFIX)) or getListIn(q)!=getListIn(qp)):
          style = ", style = dashed"
          switch = 1
        if lbl.endswith(RNA_EDGE_SUFFIX):
          lbl = lbl[:-len(RNA_EDGE_SUFFIX)]
        if not switch:
          edgesTxt += '      %s -> %s [ label = "%s" %s];\n'%(newName[q.getName()],newName[qp.getName()],lbl,style)
        else:
          edgesTxt += '      %s -> %s [ dir=back, label = "" %s];\n'%(newName[qp.getName()],newName[q.getName()],style)
        
  # List of branches
  clustersTxt = ""
  clusterNum = 0
  for q in tree:
    l = tree[q]
    nodesStyle = ""
    chainTxt = ""
    previous = q
    for qp in l:
      style ="" 
      if isForbidden(qp): style = FORBIDDEN_STYLE
      elif isMandatory(qp): style = MANDATORY_STYLE
      elif qp.isFinal(): style = FINAL_STYLE
      else: style = NORMAL_STYLE
      nodesStyle += '\n          %s [label="%s",%s];'%(newName[qp.getName()],newLabels[qp.getName()],style)
      if qp != q:
        lbl = ""
        for t in previous.getOutTransition():
          if t.getSucc() == qp:
            lbl = t.getLabel()
            chainTxt += '\n          %s -> %s [style="invis"] ;'%(newName[previous.getName()],newName[qp.getName()])
        previous = qp
    if len(l)>1:    
      tmpTxt = """        %s
"""%(nodesStyle)
    else:
      tmpTxt = """      %s\n"""%(nodesStyle)
    clustersTxt += tmpTxt
  #print clustersTxt
  #print ranktxt
  initialTxt = ""
  if initState is not None:
    w = getListIn(initState)
    initialTxt = 'Dummy_%s [style=invis,label="%s"];  Dummy_%s->%s;'%(w,newName[initState.getName()],w,newName[initState.getName()])
  ranktxt = ""
  txt = """digraph finite_state_machine {
        rankdir = TB;
        sep = 8.0;
        nodesep = 0.3;
        splines= 1;
%s
%s
%s
%s
}
"""%(clustersTxt,ranktxt,edgesTxt,initialTxt)
  tmpfilepath = outpath+".dot"
  tmpfile = open(tmpfilepath,"w")
  tmpfile.write(txt)
  tmpfile.close()
  # Dot + Neato
  os.system("%s %s | %s -n -Teps -o %s"%(DOT_BIN,tmpfilepath,NEATO_BIN,outpath))
  # Dot only
  #os.system("%s %s -Tps -o %s"%(DOT_BIN,tmpfilepath,outpath))
  #os.remove(tmpfilepath)


def injectWords(aut,root,wSet,pSetStr,typeStr,stateType,terminalSates):
  for m in wSet:
    q = root
    for j in range(len(m)):
      if q.getOutTransition(m[j]) == []:
        isFinal = False
        suffix = ""
        if j == len(m)-1:
          suffix = typeStr
          isFinal = True
        child = aut.addState("%s_%s%s"%(m[:j+1],pSetStr,suffix),isFinal,False)
        if j == len(m)-1:
          child.setData(stateType,True)
        aut.addTransition(q.getName(),child.getName(),m[j])
      q = q.getOutTransition(m[j])[0].getSucc()
    terminalSates.add(q)      
    q.setFinal(True)
 

def ahoCorasickAutomaton(forbidden,mandatory, suffix= "e"):
  aut = DFA(['a','c','g','u'])
  root = aut.addState("%s%s"%(INITIAL_PREFIX,suffix),False,True)

  (forbidden_AC_states,mandatory_AC_states) = (set(),set())

  # Initial prefix tree
  injectWords(aut,root,forbidden,suffix,FORBIDDEN_SUFFIX,"f",forbidden_AC_states)
  injectWords(aut,root,mandatory,suffix,MANDATORY_SUFFIX,"m",mandatory_AC_states)
  
  # Computing failure function 
  h = {}
  h[root] = None
  for child in root.getOutTransition():
    h[child.getSucc()] = root
  seen = set()
  def getSuccStates(trans):
    q = trans.getSucc()
    seen.add(q)
    return q
  queue = list(map(getSuccStates,root.getOutTransition()[:]))
  while len(queue)>0:
    qp = queue.pop(0)
    ql = qp.getInTransition()
    if len(ql) != 1:
      raise Exception("Error: After building prefix tree, more than one incoming transition detected for %s."%(q.getName()))
    q = ql[0].getPred()
    c = ql[0].getLabel()
    r = h[q]
    while (r is not None) and (len(r.getOutTransition(c))==0):
      r = h[r]
    if r is None:
      h[qp] = root
    else:
      h[qp] = r.getOutTransition(c)[0].getSucc()
    for t in qp.getOutTransition():
      if t.getSucc() not in seen:
        queue.append(getSuccStates(t))

  # Adding failure transitions
  for q in h:
    qp = h[q]
    if qp is not None:
      aut.addTransition(q.getName(),qp.getName(),FAILURE_TRANSITION)
  aut.addTransition(root.getName(),root.getName(),FAILURE_TRANSITION)
  return (aut,forbidden_AC_states,mandatory_AC_states)


def buildPowersets(init,inset=[],outset=[]):
  if len(init)==0:
    return [(inset[:],outset[:])]
  else:
    return buildPowersets(init[1:],inset+[init[0]],outset) +  buildPowersets(init[1:],inset,outset+[init[0]])

def labelPSet(init,inset):
  res = ""
  for i in init:
    if i in inset:
      res += "1"
    else:
      res += "0"
  return res

def formatName(q,psetLabel):
  return q.getName().replace("_O","_%s"%(psetLabel))


def getPrefix(q):
  offset = q.getName().find("_")
  if offset != -1:
    return q.getName()[:offset]
  else:
    return q.getName()

def getListIn(q):
  data = q.getName().split("_")
  if len(data)>1:  return data[1]
  else: "e"


def isMandatory(q):
  return q.getName().endswith(MANDATORY_SUFFIX)

def setForbidden(q):
  q.setData("f",True)

def isForbidden(q):
  return (q.getData("f") is not None)

def getNextACState(a,q,lbl):
  l = q.getOutTransition(lbl)
  if len(l)>0:
    t = l[0]
    return t.getSucc()
  else:
    l = q.getOutTransition(FAILURE_TRANSITION)
    t = l[0]
    qp = t.getSucc()
    if q != qp:
      return getNextACState(a,qp,lbl)
    else:
      return q

def getACState(a,w,psetlbl):
  q = a.getState("e_%s"%psetlbl)
  for l in w:
    q = getNextACState(a,q,l)
  return q
  
def RNADesignAutomaton(forbidden,mandatory):
  print((forbidden, mandatory))
  alphabet = ['A','C','G','U']
  result = DFA(alphabet)
  # Adding ground state
  ground = result.addState(GROUND_STATE_NAME,False,False)
  for b in alphabet:
    result.addTransition(ground.getName(),ground.getName(),b+RNA_EDGE_SUFFIX)
  l =  buildPowersets(mandatory)

  #Order powersets by number of satisfied mandatory motifs (ensures that destinations have already been treated)
  orderedSubsets = []
  for (lin,lout) in l:
    orderedSubsets.append((len(lin),lin,lout))
  orderedSubsets.sort(reverse=True)
  print(orderedSubsets)
  automata = {}
  nstates = {}

  (global_forbidden_states,global_mandatory_states) = (set(),set())

  for (c,lin, lout) in orderedSubsets:
    psetLabel = labelPSet(mandatory,lin)
    (a,local_forb,local_mand) = ahoCorasickAutomaton(forbidden,lout,psetLabel)
    global_forbidden_states.update(local_forb)
    global_mandatory_states.update(local_mand)
    automata[psetLabel] = a
    for q in a.getStates():
      oldName = formatName(q,psetLabel)
      allMandatory = (len(lout)==0)
      noneMandatory = (len(lin)==0)
      if q in global_mandatory_states:
        newMotif = getPrefix(q)
        nl = lin+[newMotif]
        newPSetLbl = labelPSet(mandatory,nl)
        qp = getACState(automata[newPSetLbl],newMotif,newPSetLbl)
        qpLbl = formatName(qp,newPSetLbl)
        print(("Mand:",oldName,"->",qpLbl)) #nstates[qpLbl].getName()
        nstates[oldName] = nstates[qpLbl]
      #elif q in global_forbidden_states:
      #  print "Forb:",oldName,"->",ground.getName()
      #  nstates[oldName] = ground
      else:
        qp = result.addState(oldName,allMandatory,q.isInitial() and noneMandatory)
        if isForbidden(q):
          qp.setData("f",True)
        nstates[oldName] = qp
    # Copying transitions
    for q in a.getStates():
      oldName = formatName(q,psetLabel)
      #print newName
      if isMandatory(q):
        # Already rerouted to augmented set
        pass
      #elif isForbidden(q):
        # Already rerouted to ground "garbage" state
      #  pass
      else:
        for lbl in alphabet:
          l = q.getOutTransition(lbl)
          suff = RNA_EDGE_SUFFIX
          if (len(l)==1):
            suff = ""
          q1 = nstates[oldName]
          if q.getName() == q1.getName():
            qp = getNextACState(a,q,lbl)
            q2 = nstates[formatName(qp,psetLabel)]          
            print(("%s[%s]->%s[%s]"%(q1.getName(),q1.getData("f"),q2.getName(),q2.getData("f"))))
            result.addTransition(q1.getName(),q2.getName(),lbl+suff)
  return minimizeAutomaton(result)

def minimizeAutomaton(aut):
  state2Class = {}
  stateTrans = {}
  nbClasses = 0
  classes = {0:set(),1:set(),2:set(),3:set()}
  alpha = aut.getAlphabet()
  finals = aut.getFinals()
  for q in aut.getStates():
    fin = (q in finals)
    forb = isForbidden(q)
    if not fin and not forb:
      nbClass = 0
    elif fin and not forb:
      nbClass = 1
    elif not fin and forb:
      nbClass = 2
    elif fin and forb:
      nbClass = 3
    state2Class[q] = nbClass
    classes[nbClass].add(q)
    stateTrans[q] = {}
    for c in alpha:
      t = q.getOutTransition(c)
      if len(t)!= 1:
        raise Exception("More or less than one out transition for letter [%s] in state [%s]!"%(c,q))
      stateTrans[q][c] = t[0].getSucc()
      
  while (nbClasses != len(classes)):
    nbClasses = len(classes)
    classes = {}
    for q in state2Class:
      cl = state2Class[q]
      clout = tuple([state2Class[stateTrans[q][c]] for c in alpha])
      if (cl,clout) not in classes:
        classes[(cl,clout)] = set()
      classes[(cl,clout)].add(q)
    #print classes.keys()
    nextClass = 0
    for trans in classes:
      for q in classes[trans]:
        state2Class[q] = nextClass
      nextClass += 1
    #print state2Class.values()
  classes = {}
  for q in state2Class:
    cl = state2Class[q]
    clout = tuple([state2Class[stateTrans[q][c]] for c in alpha])
    if (cl,clout) not in classes:
      classes[(cl,clout)] = set()
    classes[(cl,clout)].add(q)
      
  naut = DFA(aut.getAlphabet())
  for (q,clout) in classes:
    isInitial = False
    isFinal = True
    isForb = True
    for q2 in classes[(q,clout)]:
      if q2.isInitial():
        isInitial = True
      if not q2.isFinal():
        isFinal = False
      if not isForbidden(q2):
        isForb = False
    nq = naut.addState("%s"%q,isFinal,isInitial)
    if isForb:
      nq.setData("f",True) 
    
  for (q,clout) in classes:
    #print q,clout
    for i in range(len(alpha)):
      q2 = clout[i]
      c =  alpha[i]
      #print "%s->[%s] %s"%(q,c,q2)
      naut.addTransition("%s"%q,"%s"%q2,c)
  return naut


def injectStates(aut,tree,father,prefix=""):
  for c in tree:
    child = aut.addState(prefix+c,True)
    aut.addTransition(father.getName(),child.getName(),c)
    injectStates(aut,tree[c],child,prefix+c)
    
    

def buildAutomaton(forbidden,mandatory):
  res = {}
  for w in forbidden:
    tmp = res
    for c in w:
      if c not in tmp:
        tmp[c] = {}
      tmp = tmp[c]
    tmp['.'] = {}
  for w in mandatory:
    tmp = res
    for c in w:
      if c not in tmp:
        tmp[c] = {}
      tmp = tmp[c]
    tmp['*'] = {}
  #printTree(res)
  aut = DFA(['a','c','g','u'])
  base = aut.addState("",True,True)
  injectStates(aut,res,base)
  return aut

def printTree(t,depth=0):
  for c in t:
    print(("%s->%s"%(" "*depth,c)))
    printTree(t[c],depth+1)
      

if __name__ == '__main__':
  forbidden = []
  mandatory = []
  mode = 1 

  i=1
  while i<len(sys.argv):
    if sys.argv[i] == "-f":
      mode = 1
    elif sys.argv[i] == "-m":
      mode = 2
    else:
      if mode == 1:
        forbidden.append(sys.argv[i])
      else:
        mandatory.append(sys.argv[i])
    i += 1
  #print mandatory
  #print forbidden
  #aut = buildAutomaton(forbidden,[])
  aut = ahoCorasickAutomaton(forbidden,[])
  #print getTree(aut)
  #print aut
  drawAutomaton(aut,"out.eps")
  #res = RNADesignAutomaton(forbidden,mandatory)
  #drawACAutomaton(res,"out2.ps")
  #print res
